#!/bin/bash
# Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  * Neither the name of NVIDIA CORPORATION nor the names of its
#    contributors may be used to endorse or promote products derived
#    from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

############################################################################
## This script generates the model repository needed by some of the
## tritonserver CI tests. Generating these models requires using
## the PyTorch container.
##
## 1. Update TENSORRT_IMAGE and PYTORCH_IMAGE to
## match what is being used by the tritonserver release being
## tested.
##
## 2. Set CUDA_DEVICE to the ID of the CUDA device present on the
## local system that you want to target for the generated models.
##
## 3. Run this script to create /tmp/qa_model_repository,
## /tmp/qa_variable_model_repository, /tmp/qa_shapetensor_model_repository
## /tmp/qa_identity_model_repository, /tmp/qa_identity_big_model_repository
## /tmp/qa_reshape_model_repository, /tmp/qa_noshape_model_repository,
## /tmp/qa_sequence_model_repository, /tmp/qa_ensemble_model_repository
## /tmp/qa_dyna_sequence_model_repository, and
## /tmp/qa_variable_sequence_model_repository directories containing
## all the models needed for CI testing.
##
############################################################################

TRITON_VERSION=${TRITON_VERSION:=25.09}

# ONNX. Use ONNX_OPSET 0 to use the default for ONNX version
ONNX_VERSION=1.16.1
ONNX_OPSET=0

# OPENVINO version
OPENVINO_VERSION=2024.5.0

UBUNTU_IMAGE=${UBUNTU_IMAGE:=ubuntu:22.04}
PYTORCH_IMAGE=${PYTORCH_IMAGE:=nvcr.io/nvidia/pytorch:$TRITON_VERSION-py3}
TENSORRT_IMAGE=${TENSORRT_IMAGE:=nvcr.io/nvidia/tensorrt:$TRITON_VERSION-py3}
CUDA_DEVICE=${NV_GPU:=0}

DOCKER_GPU_ARGS=${DOCKER_GPU_ARGS:-$([[ $RUNNER_GPUS =~ ^[0-9] ]] && eval $NV_DOCKER_ARGS || echo "--gpus device=$CUDA_DEVICE" )}
MODEL_TYPE=${MODEL_TYPE:-""}

############################################################################
# Check if Docker volume exists
############################################################################
CI_JOB_ID=${CI_JOB_ID:=$(date +%Y%m%d_%H%M)}
DOCKER_VOLUME=${DOCKER_VOLUME:=volume.gen_qa_model_repository.${CI_JOB_ID}}
RUNNER_ID=${RUNNER_ID:=0}
PROJECT_NAME=${CI_PROJECT_NAME:=tritonserver}
DOCKER_VOLUME_CONTAINER=${DOCKER_VOLUME}.gen_qa_model_repository.${CI_JOB_ID}

if ! docker volume inspect $DOCKER_VOLUME > /dev/null 2>&1; then
    echo -e "\033[34m[ INFO ] - Docker volume $DOCKER_VOLUME does not exist. Creating... \033[0m "
    docker volume create $DOCKER_VOLUME --label RUNNER_ID=$RUNNER_ID --label PROJECT_NAME=$PROJECT_NAME
    docker volume inspect $DOCKER_VOLUME
else
    echo -e "\033[34m[ INFO ] - Docker volume in use: $DOCKER_VOLUME \033[0m "
    docker volume inspect $DOCKER_VOLUME
fi


docker run \
    --rm \
    --label RUNNER_ID=$RUNNER_ID \
    --label PROJECT_NAME=$PROJECT_NAME \
    -v $DOCKER_VOLUME:/mnt \
    -w /mnt/$CI_JOB_ID \
    $UBUNTU_IMAGE \
    mkdir -p gen_srcdir ${TRITON_VERSION}

docker create \
    --label RUNNER_ID=$RUNNER_ID \
    --label PROJECT_NAME=$PROJECT_NAME \
    --name $DOCKER_VOLUME_CONTAINER \
    -v $DOCKER_VOLUME:/mnt \
    -w /mnt/$CI_JOB_ID \
    $UBUNTU_IMAGE

docker cp . $DOCKER_VOLUME_CONTAINER:/mnt/$CI_JOB_ID/gen_srcdir

VOLUME_BUILD_DIR=${VOLUME_BUILD_DIR:=/mnt/$CI_JOB_ID}
VOLUME_SRCDIR=${VOLUME_SRCDIR:=$VOLUME_BUILD_DIR/gen_srcdir}
VOLUME_DESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_model_repository
VOLUME_VARDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_variable_model_repository
VOLUME_IDENTITYDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_identity_model_repository
VOLUME_IDENTITYBIGDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_identity_big_model_repository
VOLUME_SHAPEDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_shapetensor_model_repository
VOLUME_RESHAPEDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_reshape_model_repository
VOLUME_SEQDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_sequence_model_repository
VOLUME_DYNASEQDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_dyna_sequence_model_repository
VOLUME_DYNASEQIMPLICITDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_dyna_sequence_implicit_model_repository
VOLUME_VARSEQDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_variable_sequence_model_repository
VOLUME_ENSEMBLEDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_ensemble_model_repository
VOLUME_NOSHAPEDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_noshape_model_repository
VOLUME_PLGDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_trt_plugin_model_repository
VOLUME_RAGGEDDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_ragged_model_repository
VOLUME_FORMATDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_trt_format_model_repository
VOLUME_DATADEPENDENTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_trt_data_dependent_model_repository
VOLUME_IMPLICITSEQDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_sequence_implicit_model_repository
VOLUME_VARIMPLICITSEQDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_variable_sequence_implicit_model_repository
VOLUME_INITIALSTATEIMPLICITSEQDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_sequence_initial_state_implicit_model_repository
VOLUME_VARINITIALSTATEIMPLICITSEQDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_variable_sequence_initial_state_implicit_model_repository
VOLUME_TORCHTRTDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/torchtrt_model_store
VOLUME_SCALARMODELSDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_scalar_models
VOLUME_IMAGEMODELSDESTDIR=$VOLUME_BUILD_DIR/$TRITON_VERSION/qa_dynamic_batch_image_model_repository


docker run \
    --rm \
    --label RUNNER_ID=$RUNNER_ID \
    --label PROJECT_NAME=$PROJECT_NAME \
    -v $DOCKER_VOLUME:/mnt \
    -w /mnt/$CI_JOB_ID \
    $UBUNTU_IMAGE \
    mkdir -p \
        $VOLUME_BUILD_DIR \
        $VOLUME_SRCDIR \
        $VOLUME_DESTDIR \
        $VOLUME_VARDESTDIR \
        $VOLUME_IDENTITYDESTDIR \
        $VOLUME_SIGDEFDESTDIR \
        $VOLUME_IDENTITYBIGDESTDIR \
        $VOLUME_TFPARAMETERSDESTDIR \
        $VOLUME_SHAPEDESTDIR \
        $VOLUME_RESHAPEDESTDIR \
        $VOLUME_SEQDESTDIR \
        $VOLUME_DYNASEQDESTDIR \
        $VOLUME_DYNASEQIMPLICITDESTDIR \
        $VOLUME_VARSEQDESTDIR \
        $VOLUME_ENSEMBLEDESTDIR \
        $VOLUME_NOSHAPEDESTDIR \
        $VOLUME_PLGDESTDIR \
        $VOLUME_RAGGEDDESTDIR \
        $VOLUME_FORMATDESTDIR \
        $VOLUME_DATADEPENDENTDIR \
        $VOLUME_IMPLICITSEQDESTDIR \
        $VOLUME_VARIMPLICITSEQDESTDIR \
        $VOLUME_INITIALSTATEIMPLICITSEQDESTDIR \
        $VOLUME_VARINITIALSTATEIMPLICITSEQDESTDIR \
        $VOLUME_TORCHTRTDESTDIR \
        $VOLUME_SCALARMODELSDESTDIR \
        $VOLUME_IMAGEMODELSDESTDIR

ONNXSCRIPT=gen.ONNXRuntime.gen_qa_model_repository.cmds
OPENVINOSCRIPT=gen.OpenVINO.gen_qa_model_repository.cmds
TORCHSCRIPT=gen.PyTorch.gen_qa_model_repository.cmds
TRTSCRIPT=gen.TensorRT.gen_qa_model_repository.cmds

# OPENVINO
#
# OpenVINO is not available on ARM so skip
if [[ "aarch64" != $(uname -m) ]] ; then

cat > $OPENVINOSCRIPT <<EOF
#!/bin/bash
# Make all generated files accessible outside of container
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
set -e
set -x
export DEBIAN_FRONTEND=noninteractive
apt-get update && \
    apt-get install -y --no-install-recommends build-essential cmake libprotobuf-dev \
            protobuf-compiler python3 python3-dev python3-pip wget gnupg2 \
            software-properties-common
ln -s /usr/bin/python3 /usr/bin/python

pip3 install  "numpy<=1.23.5" setuptools

pip3 install openvino==$OPENVINO_VERSION

# Since variable shape tensors are not allowed, identity models may fail to generate.
# TODO Add variable size tensor models after DLIS-2827 adds support for variable shape tensors.
# TODO Add sequence models after DLIS-2864 adds support for sequence/control inputs.
python3 $VOLUME_SRCDIR/gen_qa_models.py --openvino --models_dir=$VOLUME_DESTDIR
chmod -R 777 $VOLUME_DESTDIR
# python3 $VOLUME_SRCDIR/gen_qa_identity_models.py --openvino --models_dir=$VOLUME_IDENTITYDESTDIR
# chmod -R 777 $VOLUME_IDENTITYDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_reshape_models.py --openvino --models_dir=$VOLUME_RESHAPEDESTDIR
chmod -R 777 $VOLUME_RESHAPEDESTDIR
# python3 $VOLUME_SRCDIR/gen_qa_sequence_models.py --openvino --models_dir=$VOLUME_SEQDESTDIR
# chmod -R 777 $SVOLUME_EQDESTDIR
# python3 $VOLUME_SRCDIR/gen_qa_dyna_sequence_models.py --openvino --models_dir=$VOLUME_DYNASEQDESTDIR
# chmod -R 777 $VOLUME_DYNASEQDESTDIR
EOF

chmod a+x $OPENVINOSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker cp $OPENVINOSCRIPT $DOCKER_VOLUME_CONTAINER:$VOLUME_SRCDIR

docker pull $UBUNTU_IMAGE

echo -e "\033[34m[ INFO ] - Running:  $OPENVINOSCRIPT \033[0m "

docker run \
    --rm \
    --label RUNNER_ID=$RUNNER_ID \
    --label PROJECT_NAME=$PROJECT_NAME \
    $DOCKER_GPU_ARGS \
    -v $DOCKER_VOLUME:/mnt \
    $UBUNTU_IMAGE \
    bash -xe $VOLUME_SRCDIR/$OPENVINOSCRIPT

if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi

fi # [[ "aarch64" != $(uname -m) ]]

# ONNX
cat > $ONNXSCRIPT <<EOF
#!/bin/bash
# Make all generated files accessible outside of container
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
set -e
set -x
export DEBIAN_FRONTEND=noninteractive
apt-get update && \
        apt-get install -y --no-install-recommends build-essential cmake libprotobuf-dev \
                protobuf-compiler python3 python3-dev python3-pip
ln -s /usr/bin/python3 /usr/bin/python

pip3 install "protobuf<=3.20.1"  "numpy<=1.23.5" # TODO: Remove current line DLIS-3838
pip3 install --upgrade onnx==${ONNX_VERSION}

python3 $VOLUME_SRCDIR/gen_qa_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$VOLUME_DESTDIR
chmod -R 777 $VOLUME_DESTDIR
python3 $VOLUME_SRCDIR/gen_qa_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$VOLUME_VARDESTDIR
chmod -R 777 $VOLUME_VARDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_identity_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$VOLUME_IDENTITYDESTDIR
chmod -R 777 $VOLUME_IDENTITYDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_reshape_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$VOLUME_RESHAPEDESTDIR
chmod -R 777 $VOLUME_RESHAPEDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$VOLUME_SEQDESTDIR
chmod -R 777 $VOLUME_SEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$VOLUME_VARSEQDESTDIR
chmod -R 777 $VOLUME_VARSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_implicit_models.py --onnx --initial-state zero --onnx_opset=$ONNX_OPSET --models_dir=$VOLUME_INITIALSTATEIMPLICITSEQDESTDIR
chmod -R 777 $VOLUME_INITIALSTATEIMPLICITSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_implicit_models.py --onnx --initial-state zero --onnx_opset=$ONNX_OPSET --variable --models_dir=$VOLUME_VARINITIALSTATEIMPLICITSEQDESTDIR
chmod -R 777 $VOLUME_VARINITIALSTATEIMPLICITSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_implicit_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$VOLUME_IMPLICITSEQDESTDIR
chmod -R 777 $VOLUME_IMPLICITSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_implicit_models.py --onnx --onnx_opset=$ONNX_OPSET --variable --models_dir=$VOLUME_VARIMPLICITSEQDESTDIR
chmod -R 777 $VOLUME_VARIMPLICITSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_dyna_sequence_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$VOLUME_DYNASEQDESTDIR
chmod -R 777 $VOLUME_DYNASEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_dyna_sequence_implicit_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$VOLUME_DYNASEQIMPLICITDESTDIR
chmod -R 777 $VOLUME_DYNASEQIMPLICITDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_ragged_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$VOLUME_RAGGEDDESTDIR
chmod -R 777 $VOLUME_RAGGEDDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_ort_scalar_models.py --onnx_opset=$ONNX_OPSET --models_dir=$VOLUME_SCALARMODELSDESTDIR
chmod -R 777 $VOLUME_SCALARMODELSDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_models.py --ensemble --models_dir=$VOLUME_ENSEMBLEDESTDIR/qa_model_repository
python3 $VOLUME_SRCDIR/gen_qa_models.py --ensemble --variable --models_dir=$VOLUME_ENSEMBLEDESTDIR/qa_variable_model_repository
python3 $VOLUME_SRCDIR/gen_qa_reshape_models.py --ensemble --models_dir=$VOLUME_ENSEMBLEDESTDIR/qa_reshape_model_repository
python3 $VOLUME_SRCDIR/gen_qa_identity_models.py --ensemble --models_dir=$VOLUME_ENSEMBLEDESTDIR/qa_identity_model_repository
python3 $VOLUME_SRCDIR/gen_qa_sequence_models.py --ensemble --models_dir=$VOLUME_ENSEMBLEDESTDIR/qa_sequence_model_repository
python3 $VOLUME_SRCDIR/gen_qa_sequence_models.py --ensemble --variable --models_dir=$VOLUME_ENSEMBLEDESTDIR/qa_variable_sequence_model_repository
chmod -R 777 $VOLUME_ENSEMBLEDESTDIR
EOF

chmod a+x $ONNXSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker cp $ONNXSCRIPT $DOCKER_VOLUME_CONTAINER:$VOLUME_SRCDIR

docker pull $UBUNTU_IMAGE

echo -e "\033[34m[ INFO ] - Running:  $ONNXSCRIPT \033[0m "

docker run \
    --rm \
    --label RUNNER_ID=$RUNNER_ID \
    --label PROJECT_NAME=$PROJECT_NAME \
    $DOCKER_GPU_ARGS \
    -v $DOCKER_VOLUME:/mnt \
    $UBUNTU_IMAGE \
    bash -xe $VOLUME_SRCDIR/$ONNXSCRIPT

if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi

# PyTorch
cat > $TORCHSCRIPT <<EOF
#!/bin/bash
# Make all generated files accessible outside of container
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
pip3 install onnxscript
set -e
set -x
python3 $VOLUME_SRCDIR/gen_qa_models.py --libtorch --models_dir=$VOLUME_DESTDIR
chmod -R 777 $VOLUME_DESTDIR
python3 $VOLUME_SRCDIR/gen_qa_models.py --libtorch --variable --models_dir=$VOLUME_VARDESTDIR
chmod -R 777 $VOLUME_VARDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_identity_models.py --libtorch --models_dir=$VOLUME_IDENTITYDESTDIR
chmod -R 777 $VOLUME_IDENTITYDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_reshape_models.py --libtorch --variable --models_dir=$VOLUME_RESHAPEDESTDIR
chmod -R 777 $VOLUME_RESHAPEDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_sequence_models.py --libtorch --models_dir=$VOLUME_SEQDESTDIR
chmod -R 777 $VOLUME_SEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_sequence_models.py --libtorch --variable --models_dir=$VOLUME_VARSEQDESTDIR
chmod -R 777 $VOLUME_VARSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_implicit_models.py --libtorch --models_dir=$VOLUME_IMPLICITSEQDESTDIR
chmod -R 777 $VOLUME_IMPLICITSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_implicit_models.py --libtorch --variable --models_dir=$VOLUME_VARIMPLICITSEQDESTDIR
chmod -R 777 $VOLUME_VARIMPLICITSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_dyna_sequence_models.py --libtorch --models_dir=$VOLUME_DYNASEQDESTDIR
chmod -R 777 $VOLUME_DYNASEQDESTDIR
if [ -z "$MODEL_TYPE" ] || [ "$MODEL_TYPE" != "igpu" ]; then
  nvidia-smi --query-gpu=compute_cap | grep -qz 12.1 && echo -e '\033[33m[WARNING]\033[0m Skipping model generation for Torch TensorRT' || python3 $VOLUME_SRCDIR/gen_qa_torchtrt_models.py --models_dir=$VOLUME_TORCHTRTDESTDIR
  chmod -R 777 $VOLUME_TORCHTRTDESTDIR
fi
python3 $VOLUME_SRCDIR/gen_qa_ragged_models.py --libtorch --models_dir=$VOLUME_RAGGEDDESTDIR
chmod -R 777 $VOLUME_RAGGEDDESTDIR
# Export torchvision image models to ONNX
python3 $VOLUME_SRCDIR/gen_qa_image_models.py --resnet50 --resnet152 --vgg19 --models_dir=$VOLUME_IMAGEMODELSDESTDIR
chmod -R 777 $VOLUME_IMAGEMODELSDESTDIR
EOF

chmod a+x $TORCHSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

docker cp $TORCHSCRIPT $DOCKER_VOLUME_CONTAINER:$VOLUME_SRCDIR

docker pull $PYTORCH_IMAGE

echo -e "\033[34m[ INFO ] - Running:  $TORCHSCRIPT \033[0m "

docker run \
    --rm \
    --label RUNNER_ID=$RUNNER_ID \
    --label PROJECT_NAME=$PROJECT_NAME \
    $DOCKER_GPU_ARGS \
    -v $DOCKER_VOLUME:/mnt \
    $PYTORCH_IMAGE \
    bash -xe $VOLUME_SRCDIR/$TORCHSCRIPT

if [ $? -ne 0 ]; then
    echo -e "Failed"
    exit 1
fi

# TensorRT

docker pull ${TENSORRT_IMAGE}
TENSORRT_VERSION="$(docker inspect ${TENSORRT_IMAGE} --format '{{index .Config.Labels "com.nvidia.tensorrt.version"}}' | cut -d . -f -2)"

cat > $TRTSCRIPT <<EOF
#!/bin/bash
# Make all generated files accessible outside of container
umask 0000
nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true
nvidia-smi || true
set -e
set -x
dpkg -l | grep TensorRT
export TRT_SUPPRESS_DEPRECATION_WARNINGS=1
# Models using shape tensor i/o
python3 $VOLUME_SRCDIR/gen_qa_identity_models.py --tensorrt-shape-io --models_dir=$VOLUME_SHAPEDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_sequence_models.py --tensorrt-shape-io --models_dir=$VOLUME_SHAPEDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_dyna_sequence_models.py --tensorrt-shape-io --models_dir=$VOLUME_SHAPEDESTDIR
chmod -R 777 $VOLUME_SHAPEDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_models.py --tensorrt --models_dir=$VOLUME_DESTDIR
chmod -R 777 $VOLUME_DESTDIR
python3 $VOLUME_SRCDIR/gen_qa_models.py --tensorrt --variable --models_dir=$VOLUME_VARDESTDIR
chmod -R 777 $VOLUME_VARDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_identity_models.py --tensorrt --models_dir=$VOLUME_IDENTITYDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_identity_models.py --tensorrt-compat --models_dir=$VOLUME_IDENTITYDESTDIR
chmod -R 777 $VOLUME_IDENTITYDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_identity_models.py --tensorrt-big --models_dir=$VOLUME_IDENTITYBIGDESTDIR
chmod -R 777 $VOLUME_IDENTITYBIGDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_reshape_models.py --tensorrt --variable --models_dir=$VOLUME_RESHAPEDESTDIR
chmod -R 777 $VOLUME_RESHAPEDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_sequence_models.py --tensorrt --models_dir=$VOLUME_SEQDESTDIR
chmod -R 777 $VOLUME_SEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_implicit_models.py --tensorrt --models_dir=$VOLUME_IMPLICITSEQDESTDIR
chmod -R 777 $VOLUME_IMPLICITSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_implicit_models.py --tensorrt --variable --models_dir=$VOLUME_VARIMPLICITSEQDESTDIR
chmod -R 777 $VOLUME_VARIMPLICITSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_dyna_sequence_models.py --tensorrt --models_dir=$VOLUME_DYNASEQDESTDIR
chmod -R 777 $VOLUME_DYNASEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_sequence_models.py --tensorrt --variable --models_dir=$VOLUME_VARSEQDESTDIR
chmod -R 777 $VOLUME_VARSEQDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_dyna_sequence_implicit_models.py --tensorrt --models_dir=$VOLUME_DYNASEQIMPLICITDESTDIR
chmod -R 777 $VOLUME_DYNASEQIMPLICITDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_ragged_models.py --tensorrt --models_dir=$VOLUME_RAGGEDDESTDIR
chmod -R 777 $VOLUME_RAGGEDDESTDIR
python3 $VOLUME_SRCDIR/gen_qa_trt_format_models.py --models_dir=$VOLUME_FORMATDESTDIR
chmod -R 777 $VOLUME_FORMATDESTDIR
nvidia-smi --query-gpu=compute_cap | grep -qz 11.0 && echo -e '\033[33m[WARNING]\033[0m Skipping model generation for data dependent shape' || python3 $VOLUME_SRCDIR/gen_qa_trt_data_dependent_shape.py --models_dir=$VOLUME_DATADEPENDENTDIR
chmod -R 777 $VOLUME_DATADEPENDENTDIR
# Make shared library for custom Hardmax plugin.
if [ -d "/usr/src/tensorrt" ]; then
    cd /usr/src/tensorrt/samples/python/onnx_custom_plugin
else
    git clone -b release/${TENSORRT_VERSION} https://github.com/NVIDIA/TensorRT.git
    cd /workspace/TensorRT/samples/python/onnx_custom_plugin
fi
rm -rf build && mkdir build && \
cd build && cmake .. && make -j && cp libcustomHardmaxPlugin.so $VOLUME_PLGDESTDIR/.
LD_PRELOAD=$VOLUME_PLGDESTDIR/libcustomHardmaxPlugin.so python3 $VOLUME_SRCDIR/gen_qa_trt_plugin_models.py --models_dir=$VOLUME_PLGDESTDIR
chmod -R 777 $VOLUME_PLGDESTDIR
EOF

chmod a+x $TRTSCRIPT
if [ $? -ne 0 ]; then
    echo -e "Failed: chmod"
    exit 1
fi

if [ "$MODEL_TYPE" != "igpu" ] ; then
  docker cp $TRTSCRIPT $DOCKER_VOLUME_CONTAINER:$VOLUME_SRCDIR
  docker pull $TENSORRT_IMAGE

  echo -e "\033[34m[ INFO ] - Running:  $TRTSCRIPT \033[0m "

  docker run \
    --rm \
    --label RUNNER_ID=$RUNNER_ID \
    --label PROJECT_NAME=$PROJECT_NAME \
    $DOCKER_GPU_ARGS \
    -v $DOCKER_VOLUME:/mnt \
    -e TRT_VERBOSE \
    $TENSORRT_IMAGE \
    bash -xe $VOLUME_SRCDIR/$TRTSCRIPT

  if [ $? -ne 0 ]; then
      echo -e "Failed"
      exit 1
  fi

fi

if [ -z $CI ] ; then
    echo -e "\033[34m[ INFO ] - Copying generated models to /tmp/ \033[0m "
    docker cp $DOCKER_VOLUME_CONTAINER:$VOLUME_BUILD_DIR/$TRITON_VERSION /tmp/
    echo -e "\033[34m[ INFO ] - Removing Docker container $DOCKER_VOLUME_CONTAINER \033[0m "
    docker rm -f $(docker ps -a --filter volume=$DOCKER_VOLUME --format '{{ .ID }}')
    echo -e "\033[34m[ INFO ] - Removing Docker volume $DOCKER_VOLUME \033[0m "
    docker volume rm $DOCKER_VOLUME
fi
